import pandas as pd
import time
from efficient_apriori import apriori

data = pd.read_csv('./Market_Basket_Optimisation.csv', header=None)
data1 = pd.read_csv('./Market_Basket_Optimisation2.csv', header=None)
#print(data.head(5))

def rule1():
    start = time.time()
    transactions = []
    for row in data.iterrows():
        transactions.append((set(a_ for a_ in row[1].values.tolist() if a_ == a_)))

    #print(transactions)
    # 挖掘频繁项集和关联规则
    itemsets, rules = apriori(transactions, min_support=0.02, min_confidence=0.4)
    print('频繁项集：', itemsets)
    print('关联规则：', rules)

    end = time.time()
    print("用时：", end - start)

# 采用mlxtend.frequent_patterns工具包
def rule2():
    from mlxtend.frequent_patterns import apriori
    from mlxtend.frequent_patterns import association_rules

    pd.options.display.max_columns=100
    start = time.time()
    hot_encoded_df= data1[0].str.get_dummies('|')
    frequent_itemsets = apriori(hot_encoded_df, min_support=0.05, use_colnames=True)
    rules = association_rules(frequent_itemsets, metric="lift", min_threshold=1)
    print("频繁项集：", frequent_itemsets)
    #print("关联规则：", rules[ (rules['lift'] >= 1) & (rules['confidence'] >= 0.2) ])
    rules = rules.sort_values(by="lift", ascending=False)
    print("关联规则：",rules)
    #print(rules['confidence'])
    end = time.time()
    print("用时：", end-start)


rule1()
print('-'*100)
rule2()